package com.demo.gateway.filter;

import com.demo.common.base.constant.LogConstant;
import com.demo.gateway.constant.SysConstant;
import io.netty.buffer.EmptyByteBuf;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.time.LocalDateTime;
import java.util.UUID;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * gateway全局访问日志
 *
 * @author DayDayUp
 * create on 2023-05-19
 **/
@Slf4j
@Component
public class AccessLogGlobalFilter implements GlobalFilter, Ordered {
    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {

        return cacheRequestBody(exchange, serverHttpRequest -> {
            // don't mutate and build if same request object
            exchange.getAttributes().put(SysConstant.START_TIME, System.currentTimeMillis());
            if (serverHttpRequest == exchange.getRequest()) {
                return chain.filter(exchange);
            }
            return chain.filter(exchange.mutate().request(serverHttpRequest).build());
        });
    }

    @Override
    public int getOrder() {
        return -1000;
    }

    //---------------------------------------------- private ---------------------------------------------------------

    private void logInfo(ServerHttpRequest request, String body) {
        HttpHeaders headers = request.getHeaders();
        //设置为可修改的 或者
//        exchange.getRequest().mutate().header(HttpHeaders.AUTHORIZATION,authorization).build();
        headers= HttpHeaders.writableHttpHeaders(headers);

        String uuid = UUID.randomUUID().toString();
        headers.set(LogConstant.TRACE_ID,uuid);
        String uri = request.getPath().value();
        String params = request.getQueryParams().toString();
        String method = request.getMethodValue();
        String ip = request.getRemoteAddress().toString();
        String newHeaders = request.getHeaders().entrySet()
                .stream()
                .map(entry -> "    " + entry.getKey() + ": [" + String.join(";", entry.getValue()) + "]")
                .collect(Collectors.joining("\n"));
        LocalDateTime localDateTime = LocalDateTime.now();

        StringBuilder sb = new StringBuilder();
        sb.append("\n==================================[API_CALL]==================================\n");
        sb.append("uri : ").append(uri).append("\n");
        sb.append("method : ").append(method).append("\n");
        sb.append("ip : ").append(ip).append("\n");
        sb.append("params : ").append(params).append("\n");
        sb.append("body : ").append(body).append("\n");
        sb.append("accessDate : ").append(localDateTime).append("\n");
        sb.append("headers : { \n").append(newHeaders).append("  \n}\n");
        sb.append("==============================================================================\n");
        log.info(String.valueOf(sb));
    }

    private Mono<Void> cacheRequestBody(ServerWebExchange exchange, Function<ServerHttpRequest, Mono<Void>> function) {
        ServerHttpResponse response = exchange.getResponse();
        ServerHttpRequest request = exchange.getRequest();
        NettyDataBufferFactory factory = (NettyDataBufferFactory) response.bufferFactory();
        // Join all the DataBuffers so we have a single DataBuffer for the body
        return DataBufferUtils.join(exchange.getRequest().getBody())
                .defaultIfEmpty(factory.wrap(new EmptyByteBuf(factory.getByteBufAllocator())))
                .map((dataBuffer) -> {
                    byte[] bytes = new byte[dataBuffer.readableByteCount()];
                    dataBuffer.read(bytes);
                    String bodyString = new String(bytes, StandardCharsets.UTF_8);
                    logInfo(request, bodyString);
                    // 这里下面的代码我原先没写，后续的转发直接失效，因为body数据被拿出来了
                    Flux<DataBuffer> cachedFlux = Flux.defer(() -> {
                        DataBuffer buffer = exchange.getResponse().bufferFactory()
                                .wrap(bytes);
                        return Mono.just(buffer);
                    });

                    return (ServerHttpRequest) new ServerHttpRequestDecorator(
                            exchange.getRequest()) {
                        @Override
                        public Flux<DataBuffer> getBody() {
                            return cachedFlux;
                        }
                    };
                }).switchIfEmpty(Mono.just(exchange.getRequest())).flatMap(function);
    }
}
